"""
工具类
"""
import lzma
import os
import tempfile
from os import path
import lib.config
import pyzipper


# 计算数字长度 占用2字节  最大为 65535
def num_to_bytes(num: int) -> bytes:
    return num.to_bytes(2, byteorder='big')


# str转byte
def str_to_bytes(s: str) -> bytes:
    return s.encode('utf-8')


# bytes转int
def bytes_to_int(b: bytes) -> int:
    return int.from_bytes(b, byteorder='big')


def bytes_to_str(b: bytes) -> str:
    return str(b, 'utf-8')


# 计算文件字节大小 占用8字节
def data_len_to_bytes(data: bytes) -> bytes:
    return len(data).to_bytes(8, byteorder='big')


# 获取文件名
def get_file_msg(p):
    return path.basename(p)


# 获取文件后缀名
def get_file_ext(p):
    return path.splitext(p)[1]


# 压缩文件
def save_zip_file(save_files_dir, save_files, out_file_name, password=None):
    temp_zip_file_name = None
    try:
        # 创建临时文件
        with tempfile.NamedTemporaryFile(delete=False) as temp_file:
            temp_zip_file_name = temp_file.name
            with pyzipper.AESZipFile(temp_zip_file_name, 'w') as archive:
                if password is not None:
                    archive.setpassword(str_to_bytes(password))
                    archive.setencryption(pyzipper.WZ_AES)
                for save_file in save_files:
                    archive.write(path.join(save_files_dir, save_file), arcname=path.join("/", get_file_msg(save_file)))

        # 读取临时文件数据
        with open(temp_zip_file_name, "rb") as temp_file:
            temp_data = temp_file.read()

        # 将标志和数据合并
        final_data = lib.config.flag + temp_data
        # 写入最终文件
        with open(out_file_name, "wb") as final_file:
            final_file.write(final_data)

    finally:
        # 删除临时文件
        if temp_zip_file_name and os.path.exists(temp_zip_file_name):
            os.remove(temp_zip_file_name)


# 解压缩文件
def extract_zip_file(zip_file_name, out_dir_path, password=None):
    # if get_file_ext(zip_file_name) != ".dp":
    #     print('Unable to recognize file extension')
    #     return False

    with open(zip_file_name, "rb") as f:
        zip_data = f.read()
        config_flag_len = len(lib.config.flag)
        data = zip_data[config_flag_len:]
        flag = zip_data[0:config_flag_len]
        flag_str = bytes_to_str(flag)
        if flag_str != bytes_to_str(lib.config.flag):
            print(f"The file [{zip_file_name}] is corrupted or not a valid DP file")
            return False

    # 使用 tempfile 创建临时文件
    with tempfile.NamedTemporaryFile(delete=False) as temp_file:
        temp_zip_file_name = temp_file.name
        temp_file.write(data)

    flag = True
    try:
        with pyzipper.AESZipFile(temp_zip_file_name, 'r') as archive:
            if password is not None:
                archive.setpassword(str_to_bytes(password))
            archive.extractall(out_dir_path)
    except RuntimeError as e:
        if "Bad password for file" in str(e):
            print('Password is incorrect')
        else:
            print(f"An unexpected error occurred: {e}")
        flag = False
    except Exception as e:
        print(f'An unexpected error occurred: {e}')
        flag = False
    finally:
        # 删除临时文件
        if temp_zip_file_name and os.path.exists(temp_zip_file_name):
            os.remove(temp_zip_file_name)

    return flag
